#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <fstream>
#include <set>
#include <map>
#include <SQLAPI.h>
#include "pin.H" 

#define MAX_MOD_NAME_LEN 256

//This global set object is used to store a single record of each basic block that is encountered during execution
set<unsigned long> BBLSet;
map<string*, pair<ADDRINT,ADDRINT>> ModuleMap;
set<string*> ExcludedModuleSet;

//These are the pin.bat options for this tool
KNOB<string> KnobOutputFile(KNOB_MODE_WRITEONCE, "pintool", "o", "bbltrace.out", "specify output file name"); //'-o' lets you specify the output file
KNOB<string> KnobBlacklistFile(KNOB_MODE_WRITEONCE, "pintool", "x", "", "specify blacklist file name"); //'-x' lets you specify a blacklist file

//Helper function that gets the base name from the full path
static string * BaseName(const string &sFullPath)
{
	size_t	iFound;
	iFound = sFullPath.rfind("\\");
	return new string(sFullPath.substr(iFound + 1));
}

// This function is called before every block
// Use the fast linkage for calls
VOID PIN_FAST_ANALYSIS_CALL SaveBBLAddr(ADDRINT aiAddr)
{
	BBLSet.insert(aiAddr);
}
    
// Pin calls this function every time a new basic block is encountered
// It inserts a call to SaveBBLAddr
VOID Trace(TRACE trace, VOID *v)
{
    // Visit every basic block  in the trace
    for (BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl))
    {
        // Insert a call to docount for every bbl, passing the number of instructions.
        // IPOINT_ANYWHERE allows Pin to schedule the call anywhere in the bbl to obtain best performance.
		BBL_InsertCall(bbl, IPOINT_ANYWHERE, AFUNPTR(SaveBBLAddr), IARG_FAST_ANALYSIS_CALL, IARG_UINT32, BBL_Address(bbl), IARG_END);
    }
}

static VOID OnImageLoad(IMG img, VOID *v)
{
	ADDRINT	aiStart = IMG_LowAddress(img);
	ADDRINT	aiEnd = IMG_HighAddress(img);
	const string &sFullPath = IMG_Name(img);
	
	string 	*sName = BaseName(sFullPath);

	ModuleMap[sName] = make_pair(aiStart,aiEnd);
}
bool IsExcluded(string* sModuleName)
{
	bool bRet = false;
	set<string*>::iterator ExcludedModuleSetIterator;
	for(ExcludedModuleSetIterator = ExcludedModuleSet.begin();ExcludedModuleSetIterator != ExcludedModuleSet.end(); ExcludedModuleSetIterator++)
	{
		if (strcmp(sModuleName->c_str(),(*ExcludedModuleSetIterator)->c_str()) == 0)
		{
			bRet = true;
			break;
		}
	}
	return bRet;
}

// This function is called when the application exits
VOID OnAppExit(INT32 code, VOID *v)
{
	char sBuffer[1024];
	set<unsigned long>::iterator BBLSetIterator;
	set<string*>::iterator ExcludedModuleSetIterator;
	map<string*, pair<ADDRINT,ADDRINT>>::iterator ModuleMapIterator;
	bool bFoundMod = false;
	
    // Write to a file since cout and cerr maybe closed by the application
    ofstream OutFile;
    OutFile.open(KnobOutputFile.Value().c_str(),ios::out|ios::app);
    OutFile.setf(ios::showbase);

	OutFile << endl << "### Loaded Modules ###" << endl;
	for(ModuleMapIterator = ModuleMap.begin(); ModuleMapIterator != ModuleMap.end(); ModuleMapIterator++)
	{
		sprintf_s((char*)&sBuffer,sizeof(sBuffer),"%s => %.8x - %.8x",(*ModuleMapIterator).first->c_str(),(*ModuleMapIterator).second.first,(*ModuleMapIterator).second.second);
		OutFile << sBuffer << endl;
	}

	OutFile << endl << "### Excluded Modules ###" << endl;
	for(ExcludedModuleSetIterator = ExcludedModuleSet.begin();ExcludedModuleSetIterator != ExcludedModuleSet.end(); ExcludedModuleSetIterator++)
	{
		OutFile << (*ExcludedModuleSetIterator)->c_str() << endl;
	}

	OutFile << endl << "### Basic Blocks Executed ###" << endl;
	//Loop through every basic block record in the set
	for (BBLSetIterator = BBLSet.begin(); BBLSetIterator != BBLSet.end(); BBLSetIterator++)
	{
		//sprintf_s((char*)&sBuffer,sizeof(sBuffer),"%.8x",*BBLSetIterator);
		//OutFile << sBuffer << " => ";

		bFoundMod = false;

		//Loop through all loaded modules until we find the one that the basic block belongs to
		for(ModuleMapIterator = ModuleMap.begin(); ModuleMapIterator != ModuleMap.end(); ModuleMapIterator++)
		{
			if((*BBLSetIterator >= (*ModuleMapIterator).second.first) && (*BBLSetIterator <= (*ModuleMapIterator).second.second))
			{
				if(!IsExcluded((*ModuleMapIterator).first))
				{
					sprintf_s((char*)&sBuffer,sizeof(sBuffer),"%s+%.8x",(*ModuleMapIterator).first->c_str(),*BBLSetIterator - (*ModuleMapIterator).second.first);
					OutFile << sBuffer << endl;
				}
				bFoundMod = true;
				break;
			}
		}
		//If we don't find a module that the basic block belongs to, that's weird, so print it out
		if(!bFoundMod)
		{
			sprintf((char*)sBuffer,"%.8x",*BBLSetIterator);
			OutFile << "UNKNOWN: " << sBuffer << endl;
		}
	}

	OutFile.close();
}


INT32 Usage()
{
    cerr << "This tool records the address of every basic block that is executed" << endl;
    cerr << endl << KNOB_BASE::StringKnobSummary() << endl;
    return -1;
}

void InitBlackList(void)
{
	ifstream InFile;
	char sFileName[MAX_MOD_NAME_LEN];
	InFile.open(KnobBlacklistFile.Value().c_str(),ifstream::in);
	if(!InFile.fail())
	{
		while(!InFile.eof())
		{
			InFile.getline((char*)&sFileName,MAX_MOD_NAME_LEN);
			if(strlen(sFileName))
			{
				string* sTemp = new string();
				sTemp->assign((char*)sFileName);
				ExcludedModuleSet.insert(sTemp);
			}
		}
		InFile.close();
	}
}
//Connect to the database
string DBConnect()
{
    SAConnection con; // create connection object
    
    try
    {
        // connect to database
        // in this example it is Oracle,
        // but can also be Sybase, Informix, DB2
        // SQLServer, InterBase, SQLBase and ODBC
        con.Connect(
            "",     // database name
            "root",   // user name
            "",   // password
            SA_MySQL_Client);

        string blah = string("We are connected!\n");
		
        // Disconnect is optional
        // autodisconnect will ocur in destructor if needed
        con.Disconnect();
		string nah = string("We are disconnected!\n");
		return blah + " " + nah;
    }
    catch(SAException &x)
    {
        // SAConnection::Rollback()
        // can also throw an exception
        // (if a network error for example),
        // we will be ready
        try
        {
            // on error rollback changes
            con.Rollback();
        }
        catch(SAException &x)
        {
			string blah = string("%s\n", (const char*)x.ErrText());
			return blah;
        }
        // print error message
		string blah = string("%s\n", (const char*)x.ErrText());
        return blah;
    }
}

//main -- initializes PIN, registers trace, image load, and termination functions, and starts the target application
int main(int argc, char * argv[])
{
	do
	{
		// Initialize pin
		if (PIN_Init(argc, argv))
		{
			return Usage();
		}
		
		//Add header to output file
		ofstream OutFile;
		OutFile.open(KnobOutputFile.Value().c_str());
		OutFile.setf(ios::showbase);
		OutFile << "#########################" << endl;
		OutFile << "### Basic Block Trace ###" << endl;
		OutFile << "#########################" << endl;
		OutFile << "Did we connect to the DB? " << DBConnect() << endl;
		OutFile.close();
		
		InitBlackList();

		// Register function to be called to instrument instructions
		TRACE_AddInstrumentFunction(Trace, 0);
		
		// Register function to be called to instrument image loads
		IMG_AddInstrumentFunction(OnImageLoad, 0);
		
		// Register Fini to be called when the application exits
		PIN_AddFiniFunction(OnAppExit, 0);
		
		// Start the program, never returns
		PIN_StartProgram();
	}while(0);

    return 0;
}